#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
DeepSeek API密钥配置工具
帮助用户快速配置DeepSeek API密钥
"""

import os
import json
from pathlib import Path

def setup_deepseek_api():
    """配置DeepSeek API密钥"""
    print("DeepSeek API密钥配置工具")
    print("=" * 30)

    print("\n获取API密钥的步骤：")
    print("1. 访问 https://platform.deepseek.com/")
    print("2. 登录或注册DeepSeek账户")
    print("3. 进入 'API Keys' 页面")
    print("4. 点击 'Create new key'")
    print("5. 输入密钥名称并创建")
    print("6. 复制生成的密钥")

    # 获取用户输入
    api_key = input("\n请粘贴您的DeepSeek API密钥: ").strip()

    if not api_key:
        print("未输入密钥，配置取消")
        return False

    # 验证密钥格式
    if not api_key.startswith('sk-'):
        print("警告: DeepSeek API密钥通常以'sk-'开头")

    # 获取基础URL（可选）
    print("\nDeepSeek API基础URL（通常不需要修改）：")
    print("默认: https://api.deepseek.com")
    base_url = input("请输入基础URL（直接回车使用默认）: ").strip()
    if not base_url:
        base_url = "https://api.deepseek.com"

    # 询问保存方式
    print("\n选择保存方式：")
    print("1. 保存到config.json文件（推荐）")
    print("2. 设置为环境变量（仅当前会话）")
    print("3. 同时使用两种方式")

    choice = input("\n请选择 (1-3): ").strip()

    project_root = Path(__file__).parent
    config_file = project_root / "config.json"

    if choice == "1":
        # 保存到配置文件
        config = {}
        if config_file.exists():
            try:
                with open(config_file, 'r', encoding='utf-8') as f:
                    config = json.load(f)
            except:
                pass

        config.update({
            'openai_api_key': api_key,
            'openai_api_base': base_url,
            'model': 'deepseek-chat',
            'temperature': 0.1
        })

        with open(config_file, 'w', encoding='utf-8') as f:
            json.dump(config, f, indent=2, ensure_ascii=False)

        print(f"\n✓ DeepSeek API密钥已保存到 {config_file}")
        print(f"✓ 模型设置为: deepseek-chat")
        print(f"✓ API基础URL: {base_url}")
        return True

    elif choice == "2":
        # 设置环境变量
        os.environ['OPENAI_API_KEY'] = api_key
        os.environ['OPENAI_API_BASE'] = base_url
        print("\n✓ DeepSeek API密钥已设置为环境变量")
        print("✓ API基础URL已设置为环境变量")
        print("注意：此设置仅在当前命令行会话中有效")
        return True

    elif choice == "3":
        # 同时使用两种方式
        # 保存到配置文件
        config = {}
        if config_file.exists():
            try:
                with open(config_file, 'r', encoding='utf-8') as f:
                    config = json.load(f)
            except:
                pass

        config.update({
            'openai_api_key': api_key,
            'openai_api_base': base_url,
            'model': 'deepseek-chat',
            'temperature': 0.1
        })

        with open(config_file, 'w', encoding='utf-8') as f:
            json.dump(config, f, indent=2, ensure_ascii=False)

        # 设置环境变量
        os.environ['OPENAI_API_KEY'] = api_key
        os.environ['OPENAI_API_BASE'] = base_url

        print(f"\n✓ DeepSeek API密钥已保存到 {config_file}")
        print("✓ DeepSeek API密钥已设置为环境变量")
        print(f"✓ 模型设置为: deepseek-chat")
        print(f"✓ API基础URL: {base_url}")
        return True

    else:
        print("无效选择，配置取消")
        return False

def check_current_config():
    """检查当前配置"""
    print("\n检查当前配置...")

    # 检查环境变量
    env_key = os.getenv("OPENAI_API_KEY")
    env_base = os.getenv("OPENAI_API_BASE")

    if env_key:
        print("✓ 环境变量中已配置API密钥")
        if env_base:
            print(f"✓ 环境变量中已配置API基础URL: {env_base}")
    else:
        print("✗ 环境变量中未配置API密钥")

    # 检查配置文件
    project_root = Path(__file__).parent
    config_file = project_root / "config.json"

    if config_file.exists():
        try:
            with open(config_file, 'r', encoding='utf-8') as f:
                config = json.load(f)
                if 'openai_api_key' in config:
                    print("✓ config.json中已配置API密钥")
                    if 'openai_api_base' in config:
                        print(f"✓ config.json中已配置API基础URL: {config['openai_api_base']}")
                    if 'model' in config:
                        print(f"✓ 配置的模型: {config['model']}")
                else:
                    print("✗ config.json中未配置API密钥")
        except:
            print("✗ config.json文件格式错误")
    else:
        print("✗ config.json文件不存在")

def test_deepseek_config():
    """测试DeepSeek配置"""
    print("\n测试DeepSeek配置...")

    try:
        import asyncio
        from langchain_openai import ChatOpenAI

        # 获取配置
        api_key = os.getenv("OPENAI_API_KEY")
        base_url = os.getenv("OPENAI_API_BASE")

        if not api_key:
            # 尝试从配置文件读取
            config_file = Path(__file__).parent / "config.json"
            if config_file.exists():
                with open(config_file, 'r', encoding='utf-8') as f:
                    config = json.load(f)
                    api_key = config.get('openai_api_key')
                    base_url = config.get('openai_api_base')

        if not api_key:
            print("✗ 未找到API密钥，无法测试")
            return False

        # 创建客户端
        client = ChatOpenAI(
            api_key=api_key,
            base_url=base_url,
            model="deepseek-chat",
            temperature=0.1
        )

        print("正在连接DeepSeek API...")

        # 测试调用
        response = client.invoke("你好，请简单介绍一下你自己")

        print("✓ DeepSeek API连接成功！")
        print(f"测试响应: {response.content[:100]}...")
        return True

    except Exception as e:
        print(f"✗ DeepSeek API连接失败: {e}")
        return False

def main():
    """主函数"""
    check_current_config()

    print("\n" + "=" * 40)
    print("1. 配置DeepSeek API密钥")
    print("2. 测试当前配置")
    print("3. 重新检查配置")
    print("4. 退出")

    choice = input("\n请选择 (1-4): ").strip()

    if choice == "1":
        if setup_deepseek_api():
            print("\n配置完成！现在可以启动智能助手了")
            print("\n是否要测试配置？(y/n): ", end="")
            test_choice = input().strip().lower()
            if test_choice in ['y', 'yes', '是']:
                test_deepseek_config()
            print("\n启动命令: python start_simple.py")
        else:
            print("\n配置失败")
    elif choice == "2":
        test_deepseek_config()
    elif choice == "3":
        main()
    elif choice == "4":
        print("再见！")
    else:
        print("无效选择")
        main()

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    except Exception as e:
        print(f"程序异常: {e}")